// @HEADER
//*********************************************************************//
//  SiYuan: A numerical PDE solver                                     //
//  Copyright (2022) YUAN Xi                                           //
//  This Software is released under the BSD 2-Clause license detailed  //
//  in the file "LICENSE" in the top-level SiYuan directory            //
//*********************************************************************//
// @HEADER

#ifndef _TEMPUS_OBSERVER_WRITE_TO_EXODUS_HPP
#define _TEMPUS_OBSERVER_WRITE_TO_EXODUS_HPP

#include "Tempus_IntegratorObserverBasic.hpp"
#include "Tempus_IntegratorBasic.hpp"
#include "Teuchos_RCP.hpp"
#include "Teuchos_Assert.hpp"

#include "Panzer_STK_Interface.hpp"
#include "Panzer_GlobalIndexer.hpp"
#include "Panzer_STK_ResponseEvaluatorFactory_SolutionWriter.hpp"

#include "TianXin_STK_Utilities.hpp"

namespace SiYuan {

  class TempusObserver_WriteToExodus : 
    public Tempus::IntegratorObserverBasic<double> {

  public:
    
    TempusObserver_WriteToExodus(const Teuchos::RCP<panzer_stk::STK_Interface>& mesh,
				   const Teuchos::RCP<const panzer::GlobalIndexer>& dof_manager,
				   const Teuchos::RCP<const panzer::LinearObjFactory<panzer::Traits> >& lof,
           const int& freq, 
           const Teuchos::RCP<panzer::ResponseLibrary<panzer::Traits> > & response_library) :
      m_mesh(mesh),
      m_dof_manager(dof_manager),
      m_lof(lof),
      output_counter(0),
      output_freq(freq),
      m_response_library(response_library)
    { 
      // get all element blocks and add them to the list
      std::vector<std::string> eBlocks;
      mesh->getElementBlockNames(eBlocks);

      panzer_stk::RespFactorySolnWriter_Builder builder;
      builder.mesh = mesh;
      m_response_library->addResponse("Main Field Output",eBlocks,builder);
    }
    
    void observeEndTimeStep(const Tempus::Integrator<double> & integrator ) final
    {
      output_counter++; 
      if( output_counter%output_freq !=0 ) return;

      Teuchos::RCP<Tempus::Stepper<double> > stepper = integrator.getStepper();
      Teuchos::RCP<const Tempus::SolutionHistory<double> > history= integrator.getSolutionHistory();
      double ctime = history->getCurrentTime();
      auto solution = history->getCurrentState()->getX();
    //  Teuchos::RCP<Tempus::SolutionState<double> > state = history->findState(ctime);
    //  Teuchos::RCP<const Thyra::VectorBase<double> > solution = stepper->getStepperX();
      
      // initialize the assembly container
      panzer::AssemblyEngineInArgs ae_inargs;
      ae_inargs.container_ = m_lof->buildLinearObjContainer();
      ae_inargs.ghostedContainer_ = m_lof->buildGhostedLinearObjContainer();
      ae_inargs.alpha = 0.0;
      ae_inargs.beta = 1.0;
      ae_inargs.evaluate_transient_terms = false;

      // initialize the ghosted container
      m_lof->initializeGhostedContainer(panzer::LinearObjContainer::X,*ae_inargs.ghostedContainer_);

      
        // initialize the x vector
        const Teuchos::RCP<panzer::ThyraObjContainer<double> > thyraContainer
          = Teuchos::rcp_dynamic_cast<panzer::ThyraObjContainer<double> >(ae_inargs.container_,true);
        thyraContainer->set_x_th(Teuchos::rcp_const_cast<Thyra::VectorBase<double> >(solution));
      

      m_response_library->addResponsesToInArgs<panzer::Traits::Residual>(ae_inargs);
      m_response_library->evaluate<panzer::Traits::Residual>(ae_inargs);
      
      m_mesh->writeToExodus(ctime);
    }

    void setOutputFreqency(int q)
    { output_freq=q; }

  /*  void observeEndIntegrator(const Tempus::Integrator<double> & integrator ) final
    {
      const Teuchos::RCP<Teuchos::FancyOStream> out = integrator.getOStream();
      Teuchos::OSTab ostab(out,0,"ScreenOutput");
      *out << "    End Subcycling ---------------------------------------------------------\n\n";
    }*/
    
  protected:

    Teuchos::RCP<panzer_stk::STK_Interface> m_mesh;
    Teuchos::RCP<const panzer::GlobalIndexer> m_dof_manager;
    Teuchos::RCP<const panzer::LinearObjFactory<panzer::Traits> > m_lof;
    Teuchos::RCP<panzer::ResponseLibrary<panzer::Traits> > m_response_library;

    int output_counter;
    int output_freq;
  };

}

#endif
